#!/usr/bin/env python3
"""
修复 ConversationSession 表的唯一约束
将 session_id 的全局唯一约束改为 (user_id, session_id) 复合唯一约束
"""

import sqlite3
import os
from pathlib import Path

def fix_session_unique_constraint():
    """修复会话表的唯一约束"""
    
    # 数据库文件路径
    db_path = Path(__file__).parent.parent / 'instance' / 'biomni.db'
    
    if not db_path.exists():
        print(f"数据库文件不存在: {db_path}")
        return False
    
    try:
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        print("开始修复 conversation_sessions 表的唯一约束...")
        
        # 1. 创建新的临时表，使用正确的约束
        cursor.execute("""
            CREATE TABLE conversation_sessions_new (
                id INTEGER PRIMARY KEY,
                user_id INTEGER NOT NULL,
                session_id VARCHAR(36) NOT NULL,
                title VARCHAR(200),
                conversation_dir VARCHAR(255) NOT NULL,
                timestamp_folder VARCHAR(50) NOT NULL,
                message_count INTEGER DEFAULT 0 NOT NULL,
                total_tokens INTEGER DEFAULT 0 NOT NULL,
                total_cost FLOAT DEFAULT 0.0 NOT NULL,
                conversation_file VARCHAR(255),
                metadata_file VARCHAR(255),
                total_size BIGINT DEFAULT 0 NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
                last_message_at DATETIME,
                FOREIGN KEY (user_id) REFERENCES users (id),
                UNIQUE (user_id, session_id)
            )
        """)
        
        # 2. 复制数据到新表
        cursor.execute("""
            INSERT INTO conversation_sessions_new 
            SELECT * FROM conversation_sessions
        """)
        
        # 3. 删除旧表
        cursor.execute("DROP TABLE conversation_sessions")
        
        # 4. 重命名新表
        cursor.execute("ALTER TABLE conversation_sessions_new RENAME TO conversation_sessions")
        
        # 5. 重新创建索引
        cursor.execute("CREATE INDEX ix_conversation_sessions_user_id ON conversation_sessions (user_id)")
        cursor.execute("CREATE INDEX ix_conversation_sessions_session_id ON conversation_sessions (session_id)")
        
        conn.commit()
        print("✅ 成功修复 conversation_sessions 表的唯一约束")
        
        # 验证修复结果
        cursor.execute("SELECT COUNT(*) FROM conversation_sessions")
        count = cursor.fetchone()[0]
        print(f"表中现有记录数: {count}")
        
        return True
        
    except Exception as e:
        print(f"❌ 修复失败: {e}")
        conn.rollback()
        return False
    finally:
        conn.close()

if __name__ == '__main__':
    success = fix_session_unique_constraint()
    if success:
        print("\n数据库迁移完成！现在可以为不同用户创建相同的 session_id 了。")
    else:
        print("\n数据库迁移失败，请检查错误信息。")